package util.crypto

import javax.crypto.{CipherInputStream, CipherOutputStream, Cipher}
import java.io.{ByteArrayInputStream, ByteArrayOutputStream, InputStream, OutputStream}
import util.StreamUtils
import java.nio.charset.Charset
import java.security.{Key, PrivateKey, PublicKey, KeyFactory}
import java.security.spec.{KeySpec, EncodedKeySpec, PKCS8EncodedKeySpec, X509EncodedKeySpec}

/**
 * Inspired by http://coding.westreicher.org/?p=23
 */
abstract class RSACoderBase(keyStream: InputStream, cipherMode: Int, usePrivate: Boolean) {

  def this(keyBytes: Array[Byte], cipherMode: Int, usePrivate: Boolean) = this(new ByteArrayInputStream(keyBytes),
    cipherMode, usePrivate)

  protected def copy(input: InputStream, output: OutputStream) {
    StreamUtils.copy(input, output)
  }

  protected def read(input: InputStream): Array[Byte] = {
    val bos = new ByteArrayOutputStream
    copy(input, bos)
    bos.close
    bos.toByteArray
  }

  // create RSA private key cipher
  private[this] lazy val cipher = {
    val cipher = Cipher.getInstance("RSA")
    val encodedKey: Array[Byte] = read(keyStream)

    val keySpec: KeySpec = if (usePrivate)
      new PKCS8EncodedKeySpec(encodedKey)
    else
      new X509EncodedKeySpec(encodedKey)

    val kf: KeyFactory = KeyFactory.getInstance("RSA")
    val pk: Key = if (usePrivate)
      kf.generatePrivate(keySpec)
    else
      kf.generatePublic(keySpec)

    cipher.init(cipherMode, pk)
    cipher
  };

  def blockCipher(bytes: Array[Byte]): Array[Byte] = {
    // string initialize 2 buffers.
    // scrambled will hold intermediate results
    var scrambled = new Array[Byte](0)

    // toReturn will hold the total result
    var toReturn = new Array[Byte](0)

    // if we encrypt we use 100 byte long blocks. Decryption requires 128 byte long blocks (because of RSA)
    val length = math.min(if (cipherMode == Cipher.ENCRYPT_MODE) 100 else 128, bytes.length)

    // another buffer. this one will hold the bytes that have to be modified in this step
    var buffer = new Array[Byte](length)

    for (i <- 0 until bytes.length) {
      // if we filled our buffer array we have our block ready for de- or encryption
      if ((i > 0) && (i % length == 0)) {
        //execute the operation
        scrambled = cipher.doFinal(buffer)
        // add the result to our total result.
        toReturn = append(toReturn, scrambled)
        // here we calculate the length of the next buffer required
        // if newLength would be longer than remaining bytes in the bytes array we shorten it.
        val newLength = if (i + length > bytes.length) {
          bytes.length - i
        } else length
        // clean the buffer array
        buffer = new Array[Byte](newLength)
      }

      // copy byte into our buffer.
      buffer(i % length) = bytes(i)
    }

    // this step is needed if we had a trailing buffer. should only happen when encrypting.
    // example: we encrypt 110 bytes. 100 bytes per run means we "forgot" the last 10 bytes. they are in the buffer array
    scrambled = cipher.doFinal(buffer)

    // final step before we can return the modified data.
    toReturn = append(toReturn,scrambled)

    toReturn
  }

  private def append(prefix: Array[Byte], suffix: Array[Byte]): Array[Byte] = {
    val toReturn: Array[Byte] = new Array[Byte](prefix.length + suffix.length);
    System.arraycopy(prefix, 0, toReturn, 0, prefix.length)
    System.arraycopy(suffix, 0, toReturn, prefix.length, suffix.length)
    toReturn
  }
}







